#include "server/server.hpp"

socket_context::~socket_context() {
    auto sock = socket;
    if (socket != INVALID_SOCKET) {
        ::closesocket(socket);
        socket = INVALID_SOCKET;
    }
    if (iocp_obj) {
        ::CloseThreadpoolIo(iocp_obj);
        iocp_obj = nullptr;
    }
    server_instance._close_hook(sock, address);
}

io_context& socket_context::create_io_context(socket_operation operation, std::vector<char> buf) {
    std::scoped_lock<std::mutex> lock(m);
    io_contexts.emplace_back(std::make_unique<io_context>(socket, operation, std::move(buf)));
    return *io_contexts.back();
}

io_context& socket_context::create_io_context(socket_operation operation) {
    std::scoped_lock<std::mutex> lock(m);
    io_contexts.emplace_back(std::make_unique<io_context>(socket, operation));
    return *io_contexts.back();
}

bool socket_context::remove_io_context(io_context* io_ctx) {
    std::scoped_lock<std::mutex> lock(m);
    auto target =
        std::find_if(io_contexts.begin(), io_contexts.end(), [&](const auto& ioctx) { return ioctx.get() == io_ctx; });
    if (target == io_contexts.end()) {
        return false;
    }
    io_contexts.erase(target);
    return true;
}

void socket_server::start() {
    win::initialize_socket_library();
    _initialize_socket();

    ldebug("websocket server is now listening.");
}

bool socket_server::send(win::socket sock, std::vector<char> buf, const std::function<void()>& callback) {
    auto socket_ctx_opt = _get_socket_context(sock);
    if (!socket_ctx_opt) {
        return false;
    }

    auto& socket_ctx = *socket_ctx_opt;

    auto& new_io_ctx = socket_ctx.create_io_context(socket_operation::SOCKET_SEND, std::move(buf));
    new_io_ctx.callback = callback;

    ::StartThreadpoolIo(socket_ctx.iocp_obj);

    win::dword bytes;
    auto send_result = ::WSASend(sock, &new_io_ctx.wsa_buffer, 1, &bytes, 0, &new_io_ctx, nullptr);
    return send_result != SOCKET_ERROR || ::GetLastError() == WSA_IO_PENDING;
}

void socket_server::stop() {
    if (!_server_context || _server_context->socket == INVALID_SOCKET) {
        return;
    }
    _clear_contexts();
    _server_context = nullptr;
    linfo("server stopped.");
}

bool socket_server::close_socket(win::socket sock) {
    // auto socket_ctx = _get_socket_context(sock);
    // return _remove_client_context(socket_ctx);
    ::shutdown(sock, 2);
    return true;
}

void socket_server::register_recv_hooks(const recv_hook_t& f) {
    _recv_hook = f;
}

void socket_server::register_close_hooks(const close_hook_t& f) {
    _close_hook = f;
}

void socket_server::register_send_hooks(const send_hook_t& f) {
    _send_hook = f;
}

void socket_server::register_accept_hooks(const accept_hook_t& f) {
    _accept_hook = f;
}

void socket_server::_handle_error(
    socket_context& socket_ctx, io_context& io_ctx, win::dword error_code, const std::string& source) {
    switch (error_code) {
    case WSA_IO_PENDING: return;
    case WSA_OPERATION_ABORTED: {
        linfo("overlapped operation aborted.");
        return;
    }
    case WAIT_TIMEOUT: {
        if (_test_socket_alive(socket_ctx.socket)) {
            lwarn("network connection timeout, trying to reconnect.");
        } else {
            lwarn("client " << nonstd::to_string(socket_ctx.address) << " wait timeout, closed unexpectly.");
            if (!_remove_client_context(&socket_ctx, &io_ctx)) {
                lwarn("unrecognized client.");
            }
        }
        return;
    }
    case ERROR_CONNECTION_ABORTED: {
        lwarn("client aborted by the local system.");
        if (!_remove_client_context(&socket_ctx, &io_ctx)) {
            lwarn("unrecognized client.");
        }
        return;
    }
    case ERROR_NETNAME_DELETED: {
        if (&socket_ctx == _server_context.get()) {
            lerror("server closed unexpectly");
        } else {
            lerror("client " << nonstd::to_string(socket_ctx.address) << " closed unexpectly.");
            if (!_remove_client_context(&socket_ctx, &io_ctx)) {
                lwarn("unrecognized client.");
            }
        }
        return;
    }
    case WSAECONNRESET: {
        lwarn("connection " << socket_ctx.socket << " was closed forcibly by the remote host");
        if (!_remove_client_context(&socket_ctx, &io_ctx)) {
            lwarn("unrecognized client");
        }
        return;
    }
    default:;
    }
    if (source.empty()) {
        lfatal("unhandled socket error " << error_code << ": " << win::get_error_message(error_code));
    } else {
        lfatal("unhandled socket error " << error_code << ": " << win::get_error_message(error_code) << " from "
                                         << source);
    }
    exit(-1);
}

void socket_server::_initialize_default_hooks() {
    _accept_hook = [](win::socket, const win::sockaddr_in& addr) {
        linfo("accepted connection from " << nonstd::to_string(addr));
    };

    _recv_hook = [](win::socket, const win::sockaddr_in& addr, std::vector<char> buf) {
        linfo("received " << buf.size() << " bytes from " << nonstd::to_string(addr));
        linfo("content: " << std::string(buf.data(), std::min(buf.size(), static_cast<size_t>(100))));
        return true;
    };

    _send_hook = [](win::socket, const win::sockaddr_in& addr, size_t bytes_transferred, size_t remain) {
        linfo(bytes_transferred << " bytes have been sent to " << nonstd::to_string(addr) << ", " << remain
                                << " bytes remained.");
    };

    _close_hook = [](win::socket, const win::sockaddr_in& addr) {
        linfo("client " << nonstd::to_string(addr) << " closed.");
    };
}

socket_context* socket_server::_get_socket_context(win::socket sock) {
    std::scoped_lock<std::mutex> lock(_connections_mutex);
    auto result = std::find_if(_client_contexts.begin(), _client_contexts.end(),
        [&](const std::unique_ptr<socket_context>& ctx) { return ctx->socket == sock; });
    if (result == _client_contexts.end()) {
        return nullptr;
    }
    return (*result).get();
}

void WINAPI socket_server::worker_thread(_Inout_ win::tp_callback_instance*, // callback instance
    _Inout_opt_ void* sock_context,                                          // socket context
    _Inout_opt_ void* overlapped,                                            // overlapped structure
    _In_ win::ulong io_result,                                               // result for io
    _In_ win::ulong_ptr bytes_transferred,                                   // bytes transferred in io
    _Inout_ win::iocp_object*                                                // io structure
) {

    // socket closed
    if (!sock_context || !overlapped) {
        linfo("worker thread closed.");
        return;
    }

    io_context& io_ctx = *static_cast<io_context*>(overlapped);
    socket_context& socket_ctx = *static_cast<socket_context*>(sock_context);

    socket_server& that = socket_ctx.server_instance;

    // handle error
    if (io_result != NO_ERROR) {
        that.handle_error(socket_ctx, io_ctx, io_result);
        return;
    }

    // socket shutdown
    if (bytes_transferred == 0 && io_ctx.operation != socket_operation::SOCKET_ACCEPT) {
        linfo("client " << nonstd::to_string(socket_ctx.address) << " shutdown.");
        if (!that._remove_client_context(&socket_ctx, &io_ctx)) {
            linfo("cannot find client.");
        }
        return;
    }

    switch (io_ctx.operation) {
    case socket_operation::SOCKET_ACCEPT:
        that._accept(socket_ctx, io_ctx, static_cast<win::dword>(bytes_transferred));
        return;
    case socket_operation::SOCKET_RECV:
        that._recv(socket_ctx, io_ctx, static_cast<win::dword>(bytes_transferred));
        return;
    case socket_operation::SOCKET_SEND:
        that._send(socket_ctx, io_ctx, static_cast<win::dword>(bytes_transferred));
        return;
    default: lfatal("unexpected socket operation."); return;
    }
}

socket_context& socket_server::_create_new_client_context(win::socket client_socket) {
    std::scoped_lock<std::mutex> lock(_connections_mutex);
    _client_contexts.emplace_back(std::make_unique<socket_context>(*this, client_socket));
    return *_client_contexts.back();
}

bool socket_server::_remove_client_context(socket_context* socket_ctx, io_context* io_ctx) {
    std::scoped_lock<std::mutex> lock(_connections_mutex);
    auto res = std::find_if(
        _client_contexts.begin(), _client_contexts.end(), [=](const auto& c) { return c.get() == socket_ctx; });
    if (res == _client_contexts.end()) {
        return false;
    }
    if (socket_ctx->io_contexts.size() > 1) {
        return socket_ctx->remove_io_context(io_ctx);
    }
    linfo("socket " << (*res)->socket << " closed");
    _client_contexts.erase(res);
    return true;
}

void socket_server::_clear_contexts() {
    std::scoped_lock<std::mutex> lock(_connections_mutex);
    _client_contexts.clear();
}

void socket_server::_initialize_socket() {
    win::sockaddr_in server_address{};

    const auto sock = ::WSASocket(AF_INET, SOCK_STREAM, 0, nullptr, 0, WSA_FLAG_OVERLAPPED);

    assert(sock != INVALID_SOCKET, "server socket initialization failed. " << win::get_last_error_message());
    ldebug("socket created.");

    server_address.sin_family = AF_INET;
    inet_pton(AF_INET, _server_ip.data(), &server_address.sin_addr.s_addr);
    server_address.sin_port = htons(_server_port);

    auto bind_result = ::bind(sock, reinterpret_cast<const sockaddr*>(&server_address), sizeof(server_address));
    assert(bind_result != SOCKET_ERROR, "server socket binding failed. " << win::get_last_error_message());
    ldebug("socket binded.");

    auto listen_result = ::listen(sock, SOMAXCONN);
    assert(listen_result != SOCKET_ERROR, "server socket listening failed" << win::get_last_error_message());
    ldebug("socket is now listening to " << _server_ip << ":" << _server_port);

    // now get acceptex function pointer
    GUID acceptex_guid = WSAID_ACCEPTEX;
    GUID get_acceptex_sock_addrs_guid = WSAID_GETACCEPTEXSOCKADDRS;
    win::dword bytes;

    auto wsa_res1 = ::WSAIoctl(sock, SIO_GET_EXTENSION_FUNCTION_POINTER, &acceptex_guid, sizeof(acceptex_guid),
        &_acceptex, sizeof(_acceptex), &bytes, nullptr, nullptr);
    assert(wsa_res1 != SOCKET_ERROR, "failed to get AcceptEx() function pointer.");

    auto wsa_res2 = ::WSAIoctl(sock, SIO_GET_EXTENSION_FUNCTION_POINTER, &get_acceptex_sock_addrs_guid,
        sizeof(get_acceptex_sock_addrs_guid), &_get_acceptex_sock_addrs, sizeof(_get_acceptex_sock_addrs), &bytes,
        nullptr, nullptr);
    assert(wsa_res2 != SOCKET_ERROR, "failed to get GetAcceptExSockAddrs() function pointer");

    _server_context = std::make_unique<socket_context>(*this, sock);
    _server_context->iocp_obj =
        ::CreateThreadpoolIo(reinterpret_cast<win::handle>(sock), worker_thread, _server_context.get(), nullptr);

    for (size_t i = 0; i < _accept_count; ++i) {
        auto& io_ctx = _server_context->create_io_context(socket_operation::SOCKET_ACCEPT);
        _post_accept(io_ctx);
    }
}

void socket_server::_post_accept(io_context& io_ctx) {
    win::dword bytes;

    io_ctx.socket = ::WSASocket(AF_INET, SOCK_STREAM, 0, nullptr, 0, WSA_FLAG_OVERLAPPED);
    assert(io_ctx.socket != INVALID_SOCKET, "socket creation failed");

    ::StartThreadpoolIo(_server_context->iocp_obj);

    bool acceptex_result = _acceptex(_server_context->socket, io_ctx.socket, io_ctx.wsa_buffer.buf, 0,
        sizeof(win::sockaddr_in) + 16, sizeof(win::sockaddr_in) + 16, &bytes, &io_ctx);

    if (!acceptex_result) {
        handle_socket_error(*_server_context, io_ctx);
    }
}

void socket_server::_post_recv(socket_context& socket_ctx, io_context& io_ctx) {
    io_ctx.clear_buffer();

    win::dword bytes = 0;
    win::dword flag = 0;

    ::StartThreadpoolIo(socket_ctx.iocp_obj);

    auto recv_result = ::WSARecv(io_ctx.socket, &io_ctx.wsa_buffer, 1, &bytes, &flag, &io_ctx, nullptr);
    if (recv_result == SOCKET_ERROR) {
        handle_socket_error(socket_ctx, io_ctx);
    }
}

void socket_server::_post_send(socket_context& socket_ctx, io_context& io_ctx) {
    win::dword bytes;

    ::StartThreadpoolIo(socket_ctx.iocp_obj);

    auto send_result = ::WSASend(io_ctx.socket, &io_ctx.wsa_buffer, 1, &bytes, 0, &io_ctx, nullptr);
    if (send_result == SOCKET_ERROR) {
        handle_socket_error(socket_ctx, io_ctx);
    }
}

void socket_server::_accept(socket_context&, io_context& io_ctx, win::dword) {
    win::sockaddr_in* client_addr = nullptr;
    win::sockaddr_in* server_addr = nullptr;
    const int addr_len = sizeof(SOCKADDR_IN);
    auto client_len = addr_len, server_len = addr_len;

    _get_acceptex_sock_addrs(io_ctx.wsa_buffer.buf, 0, addr_len + 16, addr_len + 16,
        reinterpret_cast<LPSOCKADDR*>(&server_addr), &server_len, reinterpret_cast<LPSOCKADDR*>(&client_addr),
        &client_len);

    auto& new_connection = _create_new_client_context(io_ctx.socket);

    memcpy(&new_connection.address, client_addr, addr_len);

    new_connection.iocp_obj = create_threadpool_io(io_ctx.socket, new_connection);
    assert(new_connection.iocp_obj, "CreateThreadpoolIo failed. " << win::get_last_error_message());

    auto& new_io_ctx = new_connection.create_io_context(socket_operation::SOCKET_RECV);
    _post_recv(new_connection, new_io_ctx);

    _accept_hook(io_ctx.socket, new_connection.address);

    io_ctx.clear_buffer();
    _post_accept(io_ctx);
}

void socket_server::_recv(socket_context& socket_ctx, io_context& io_ctx, win::dword bytes_transferred) {
    auto original_size = io_ctx.buffer.size();
    io_ctx.buffer.resize(bytes_transferred);
    if (!_recv_hook(socket_ctx.socket, socket_ctx.address, std::move(io_ctx.buffer))) {
        return;
    }
    auto new_size = original_size;
    if (bytes_transferred == original_size) { // expand buffer
        new_size = new_size * 3 / 2;
    }
    io_ctx.buffer.resize(new_size);
    io_ctx.wsa_buffer.len = static_cast<win::ulong>(new_size);
    io_ctx.wsa_buffer.buf = io_ctx.buffer.data();
    _post_recv(socket_ctx, io_ctx);
}

void socket_server::_send(socket_context& socket_ctx, io_context& io_ctx, win::dword bytes_transferred) {
    _send_hook(socket_ctx.socket, socket_ctx.address, bytes_transferred, io_ctx.wsa_buffer.len - bytes_transferred);

    assert(bytes_transferred <= io_ctx.wsa_buffer.len, "bytes_transferred is greater than buffer size.");

    if (bytes_transferred < io_ctx.wsa_buffer.len) {
        io_ctx.clear_base();
        io_ctx.wsa_buffer.buf += bytes_transferred;
        io_ctx.wsa_buffer.len -= bytes_transferred;

        ::StartThreadpoolIo(socket_ctx.iocp_obj);

        auto send_result = ::WSASend(socket_ctx.socket, &io_ctx.wsa_buffer, 1, &bytes_transferred, 0, &io_ctx, nullptr);
        if (send_result == SOCKET_ERROR) {
            handle_socket_error(socket_ctx, io_ctx);
        }
    } else {
        auto callback = io_ctx.callback;
        if (!socket_ctx.remove_io_context(&io_ctx)) {
            lwarn("cannot find io context");
        }
        if (callback) {
            callback();
        }
    }
}
